{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "[View the runnable example on GitHub](https://github.com/intel-analytics/BigDL/tree/main/python/nano/tutorial/notebook/inference/tensorflow/tensorflow_save_and_load_onnx.ipynb)"
      ],
      "metadata": {}
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Save and Load ONNXRuntime Model in TensorFlow\n"
      ],
      "metadata": {}
    },
    {
      "cell_type": "markdown",
      "source": [
        "This example illustrates how to save and load a TensorFlow Keras model accelerated by onnxruntime.\n",
        "In this example, we use a pretrained EfficientNetB0 model. Then, by calling `trace(model, accelerator=\"onnxruntime\"...)`, we can obtain a model accelarated by onnxruntime method provided by BigDL-Nano for inference. By calling `save(model_name, path)` , we could save the model to a folder. By calling `load(path, model_name)`, we could load the model from a folder."
      ],
      "metadata": {}
    },
    {
      "cell_type": "markdown",
      "source": [
        "To inference using Bigdl-nano InferenceOptimizer, the following packages need to be installed first. We recommend you to use [Miniconda](https://docs.conda.io/en/latest/miniconda.html) to prepare the environment and install the following packages in a conda environment.\n",
        "\n",
        "You can create a conda environment by executing:\n",
        "\n",
        "```\n",
        "# \"nano\" is conda environment name, you can use any name you like.\n",
        "conda create -n nano python=3.7 setuptools=58.0.4\n",
        "conda activate nano\n",
        "```\n",
        "\n",
        "> 📝 **Note**:\n",
        ">\n",
        "> During your installation, there may be some warnings or errors about version, just ignore them."
      ],
      "metadata": {
        "nbsphinx": "hidden"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Necessary packages for inference accelaration\n",
        "!pip install --pre --upgrade bigdl-nano[tensorflow,inference]"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "nbsphinx": "hidden"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "First, prepare model. We use an EfficientNetB0 model (`model_ft` in following code) pretrained on Imagenet dataset in this example."
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": [
        "from tensorflow.keras.applications import EfficientNetB0\n",
        "\n",
        "model_ft = EfficientNetB0(weights='imagenet')"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {}
    },
    {
      "cell_type": "markdown",
      "source": [
        "Accelerate Inference Using ONNX Runtime"
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "from bigdl.nano.tf.keras import InferenceOptimizer\n",
        "\n",
        "ort_model = InferenceOptimizer.trace(model_ft,\n",
        "                                     accelerator=\"onnxruntime\",\n",
        "                                     input_spec=tf.TensorSpec(shape=(None, 224, 224, 3))\n",
        "                                     )\n",
        "\n",
        "x = tf.random.normal(shape=(2, 224, 224, 3))\n",
        "# use the optimized model here\n",
        "y_hat = ort_model(x)\n",
        "predictions = tf.argmax(y_hat, axis=1)\n",
        "print(predictions)"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {}
    },
    {
      "cell_type": "markdown",
      "source": [
        "Save Optimized Model.\n",
        "The saved model files will be saved at \"./optimized_model_ort\" directory.\n",
        "There are 2 major files in optimized_model_ort, users only need to take \".onnx\" file for further usage:\n",
        "\n",
        "* nano_model_meta.yml: meta information of the saved model checkpoint\n",
        "* onnx_saved_model.onnx: model checkpoint for general use, describes model structure"
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": [
        "InferenceOptimizer.save(ort_model, \"./optimized_model_ort\")"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {}
    },
    {
      "cell_type": "markdown",
      "source": [
        "Load the Optimized Model"
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": [
        "loaded_model = InferenceOptimizer.load(\"./optimized_model_ort\", model_ft)"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {}
    },
    {
      "cell_type": "markdown",
      "source": [
        "Inference with the Loaded Model"
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": [
        "# use the optimized model here\n",
        "y_hat_ld = loaded_model(x)\n",
        "predictions_ld = tf.argmax(y_hat_ld, axis=1)\n",
        "print(predictions_ld)"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {}
    },
    {
      "cell_type": "markdown",
      "source": [
        "> 📚 **Related Readings**\n",
        ">\n",
        "> - [How to install BigDL-Nano](https://bigdl.readthedocs.io/en/latest/doc/Nano/Overview/install.html)\n",
        "> - [How to install BigDL-Nano in Google Colab](https://bigdl.readthedocs.io/en/latest/doc/Nano/Howto/Install/install_in_colab.html)"
      ],
      "metadata": {}
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3.7.10 ('ruonan_nano')",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.10"
    },
    "orig_nbformat": 4,
    "vscode": {
      "interpreter": {
        "hash": "d347a5dca25745bedb029e46e41f7d6c8c9b5181ecb97033e2e81a7538459254"
      }
    },
    "nteract": {
      "version": "0.28.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}